/*
 * Copyright (c) 2023. Baidu, Inc. All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *    http://www.apache.org/licenses/LICENSE-2.0
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and limitations under the License.
 */

package com.baidu.bifromq.basekv.benchmark;

import com.baidu.bifromq.basekv.store.range.ILoadTracker;
import com.baidu.bifromq.basekv.store.range.estimator.ISplitKeyEstimator;
import com.baidu.bifromq.basekv.store.range.estimator.SplitKeyEstimator;
import com.google.protobuf.ByteString;
import java.util.concurrent.ThreadLocalRandom;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.openjdk.jmh.annotations.Benchmark;
import org.openjdk.jmh.annotations.BenchmarkMode;
import org.openjdk.jmh.annotations.Group;
import org.openjdk.jmh.annotations.Level;
import org.openjdk.jmh.annotations.Mode;
import org.openjdk.jmh.annotations.Scope;
import org.openjdk.jmh.annotations.Setup;
import org.openjdk.jmh.annotations.State;
import org.openjdk.jmh.infra.Blackhole;
import org.openjdk.jmh.runner.Runner;
import org.openjdk.jmh.runner.options.Options;
import org.openjdk.jmh.runner.options.OptionsBuilder;

@Slf4j
@State(Scope.Benchmark)
public class SplitKeyEstimatorBenchmark {
    private static final int totalKeyCount = 1000000;
    private final ISplitKeyEstimator loadEstimator = new SplitKeyEstimator(5);
    private ByteString[] keys;

    @Setup(Level.Trial)
    public void setup() {
        keys = new ByteString[totalKeyCount];
        for (int i = 0; i < totalKeyCount; i++) {
            keys[i] = ByteString.copyFromUtf8("Key" + i);
        }
    }

    @Benchmark
    @Group("load")
    @BenchmarkMode(Mode.Throughput)
    public void track() {
        ILoadTracker.ILoadRecorder recorder = loadEstimator.start();
        ByteString userKey = keys[ThreadLocalRandom.current().nextInt(0, totalKeyCount)];
        recorder.record(userKey, 1);
        recorder.stop();
    }

    @Benchmark
    @Group("load")
    @BenchmarkMode(Mode.Throughput)
    public void estimate(Blackhole blackhole) {
        blackhole.consume(loadEstimator.estimate());
    }

    @SneakyThrows
    public static void main(String[] args) {
        Options opt = new OptionsBuilder()
            .include(SplitKeyEstimatorBenchmark.class.getSimpleName())
            .threads(4)
            .warmupIterations(2)
            .measurementIterations(3)
            .build();
        new Runner(opt).run();
    }
}

